package cn.itcast.spark.sql

import org.apache.spark.sql.{DataFrame, Dataset, KeyValueGroupedDataset, SparkSession}
import org.apache.spark.sql.types.{FloatType, IntegerType, StringType, StructField, StructType}
import org.junit.Test

import java.lang

class TypedTransformation {

  // 1. 创建SparkSession
  val spark = SparkSession.builder()
    .master("local[6]")
    .appName("typed")
    .getOrCreate()

  import spark.implicits._

  @Test
  def trans(): Unit = {

    // flatmap
    val ds1 = Seq("hello spark", "hello hadoop").toDS()
    ds1.flatMap(item => item.split(" ")).show()

    // map
    val ds2 = Seq(Person("zhangsan", 15), Person("lisi", 20)).toDS()
    ds2.map(person => Person(person.name, person.age * 2)).show()

    // mapPartitions
    ds2.mapPartitions(
      // iter 不能大道每个Executor的内存放不下
      // 对每个元素进行转换后，生成一个新的集合
      iter => {
        val result = iter.map(person => Person(person.name, person.age * 2))
        result
      }
    ).show()

  }

  @Test
  def trans1(): Unit = {
    val ds = spark.range(10)
    // 把一个dataset转换为另外一个dataset
    ds.transform(dataset => dataset.withColumn("doubled", 'id * 2)).show()
  }

  @Test
  def as(): Unit ={
    // 1. 读取
    val schema = StructType(
      Seq(
        StructField("name", StringType),
        StructField("age", IntegerType),
        StructField("gpa", FloatType)
      )
    )

    val df: DataFrame = spark.read
      .schema(schema)
      .option("delimiter", "\t")
      .csv("dataset/studenttab10k")

    // 2. 转换
    // 本质上 Dataset[Row].as[Student] => Dataset[Student]
    // Dataset[(String, int, float)].as[Student] => Dataset[Student]
    val ds = df.as[Student]

    // 3. 输出
    ds.show()
  }

  @Test
  def filter(): Unit = {
    val ds = Seq(Person("zhangsan", 15), Person("lisi", 20)).toDS()
    ds.filter(person => person.age > 15).show()
  }

  @Test
  def groupByKey(): Unit = {
    val ds = Seq(Person("zhangsan", 15), Person("zhangsan", 16), Person("lisi", 20)).toDS()

    val grouped: KeyValueGroupedDataset[String, Person] = ds.groupByKey(person => person.name)
    val result: Dataset[(String, Long)] = grouped.count()
    result.show()
  }

  @Test
  def split(): Unit = {
    val ds = spark.range(15)
    // randomSplit
    // 切多少分，权重多少
    val datasets: Array[Dataset[lang.Long]] = ds.randomSplit(Array(5, 2, 3))
    datasets.foreach(_.show())

    // sample
    ds.sample(withReplacement = false, fraction = 0.4).show()
  }

  @Test
  def sort(): Unit = {
    val ds = Seq(Person("zhangsan", 15), Person("zhangsan", 16), Person("lisi", 20)).toDS()
    ds.orderBy('age.desc).show()
    ds.sort('age.asc).show()
  }

  @Test
  def dropDuplicates(): Unit = {
    val ds = spark.createDataset(Seq(Person("zhangsan", 15), Person("zhangsan", 15), Person("lisi", 15)))
    ds.distinct().show()
    ds.dropDuplicates("age").show()  // 按照列名去重
  }

  @Test
  def collection(): Unit = {
    val ds1 = spark.range(1, 10)
    val ds2 = spark.range(5, 15)
    ds1.except(ds2).show()
    ds1.intersect(ds2).show()
    ds1.union(ds2).show()
    ds1.limit(5).show()
  }

}

case class Student(name: String, age: Int, gpa: Float)














